# from srcnn import init_torch_model
from srcnn import SuperResolutionNet
import os
import cv2
import numpy as np
import requests
import torch
import torch.onnx
from torch import nn

def init_torch_model():
    # # 固定upscale_factor
    # torch_model = SuperResolutionNet(upscale_factor=3)
    torch_model=SuperResolutionNet()
    state_dict = torch.load('srcnn.pth')['state_dict']

    # Adapt the checkpoint
    for old_key in list(state_dict.keys()):
        new_key = '.'.join(old_key.split('.')[1:])
        state_dict[new_key] = state_dict.pop(old_key)

    torch_model.load_state_dict(state_dict)
    torch_model.eval()
    return torch_model

model = init_torch_model()

## onnx export
x = torch.randn(1, 3, 256, 256)
#
# with torch.no_grad():
#     torch.onnx.export(
#         model,
#         x,
#         "srcnn.onnx",
#         opset_version=11,
#         input_names=['input'],
#         output_names=['output'])

##@onte 动态参数推理
with torch.no_grad():
    torch.onnx.export(
        model,
        (x, torch.tensor(3)),
        "srcnn2.onnx",
        opset_version=11,
        input_names=['input', 'factor'],
        output_names=['output'])

## onnx export
